package msfgui;

import java.awt.HeadlessException;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.util.Map;
import java.util.concurrent.locks.ReentrantLock;
import javax.swing.JOptionPane;
import javax.swing.Timer;
import javax.swing.table.DefaultTableModel;

/**
 *
 * @author scriptjunkie
 */
public class ProcessList extends MsfFrame {
	protected final Map session;
	protected final RpcConnection rpcConn;
	protected ReentrantLock lock;
	protected final DefaultTableModel model;
	protected Timer readTimer = null;

	/** Creates new form ProcessList */
	public ProcessList(final RpcConnection rpcConn, final Map session, Map sessionPopupMap) {
		super("Meterpreter remote process list");
		initComponents();
		model = new DefaultTableModel(){
			public boolean isCellEditable(int row, int col){
				return false;
			}
		};
		processTable.setModel(model);
		processTable.setShowHorizontalLines(false);
		processTable.setShowVerticalLines(false);
		this.rpcConn = rpcConn;
		this.session = session;
		this.lock = ((InteractWindow)sessionPopupMap.get(session.get("uuid"))).lock;
	}

	/** Lists the processes that are running */
	protected void listProcs() throws HeadlessException {
		try {
			rpcConn.execute("session.meterpreter_write", new Object[]{session.get("id"), Base64.encode(("ps\n").getBytes())});
		} catch (Exception ex) {
			JOptionPane.showMessageDialog(null, ex);
			if (ex.getMessage().equals("unknown session"))
				return;
		}
		readTimer = new Timer(300, new ActionListener() {
			public void actionPerformed(ActionEvent e) {
				try {
					Map received = (Map) rpcConn.execute("session.meterpreter_read", new Object[]{session.get("id")});
					if (!received.get("encoding").equals("base64")) {
						JOptionPane.showMessageDialog(null, "uhoh. encoding changed. Time to update msfgui?");
						return;
					}
					byte[] decodedBytes = Base64.decode(received.get("data").toString());
					if (decodedBytes.length == 0)
						return;
					String[] lines = new String(decodedBytes).split("\n");
					String headerRow = null;
					for (String line : lines) {
						line = line.trim();
						if(line.startsWith("PID")){
							headerRow = line;
							model.setColumnIdentifiers(line.split("\\s+"));
							while(model.getRowCount() > 0)
								model.removeRow(0);
						}
						if (line.length() < 2 || line.charAt(1) < '0' || line.charAt(1) > '9')
							continue;
						model.addRow(TableHelper.fill(line,headerRow));
					}
					readTimer.stop();
					TableHelper.fitColumnWidths(model, processTable);
				} catch(NullPointerException nex){ //junk data from previous command? Ignore.
				} catch (Exception ex) {
					ex.printStackTrace();
					if (ex.getMessage().equals("unknown session"))
						readTimer.stop();
					JOptionPane.showMessageDialog(null, ex.getMessage());
				}
			}
		});
		readTimer.start();
	}

	/** This method is called from within the constructor to
	 * initialize the form.
	 * WARNING: Do NOT modify this code. The content of this method is
	 * always regenerated by the Form Editor.
	 */
	@SuppressWarnings("unchecked")
    // <editor-fold defaultstate="collapsed" desc="Generated Code">//GEN-BEGIN:initComponents
    private void initComponents() {

        migrateButton = new javax.swing.JButton();
        killButton = new javax.swing.JButton();
        jScrollPane1 = new javax.swing.JScrollPane();
        processTable = new javax.swing.JTable();
        refreshButton = new javax.swing.JButton();

        setDefaultCloseOperation(javax.swing.WindowConstants.DISPOSE_ON_CLOSE);
        addWindowListener(new java.awt.event.WindowAdapter() {
            public void windowClosing(java.awt.event.WindowEvent evt) {
                formWindowClosing(evt);
            }
            public void windowOpened(java.awt.event.WindowEvent evt) {
                formWindowOpened(evt);
            }
        });

        org.jdesktop.application.ResourceMap resourceMap = org.jdesktop.application.Application.getInstance(msfgui.MsfguiApp.class).getContext().getResourceMap(ProcessList.class);
        migrateButton.setText(resourceMap.getString("migrateButton.text")); // NOI18N
        migrateButton.setName("migrateButton"); // NOI18N
        migrateButton.addActionListener(new java.awt.event.ActionListener() {
            public void actionPerformed(java.awt.event.ActionEvent evt) {
                migrateButtonActionPerformed(evt);
            }
        });

        killButton.setText(resourceMap.getString("killButton.text")); // NOI18N
        killButton.setName("killButton"); // NOI18N
        killButton.addActionListener(new java.awt.event.ActionListener() {
            public void actionPerformed(java.awt.event.ActionEvent evt) {
                killButtonActionPerformed(evt);
            }
        });

        jScrollPane1.setName("jScrollPane1"); // NOI18N

        processTable.setModel(new javax.swing.table.DefaultTableModel(
            new Object [][] {

            },
            new String [] {
                "Name", "PID", "User", "Session", "Path"
            }
        ) {
            Class[] types = new Class [] {
                java.lang.String.class, java.lang.String.class, java.lang.String.class, java.lang.String.class, java.lang.String.class
            };

            public Class getColumnClass(int columnIndex) {
                return types [columnIndex];
            }
        });
        processTable.setName("processTable"); // NOI18N
        jScrollPane1.setViewportView(processTable);

        refreshButton.setText(resourceMap.getString("refreshButton.text")); // NOI18N
        refreshButton.setName("refreshButton"); // NOI18N
        refreshButton.addActionListener(new java.awt.event.ActionListener() {
            public void actionPerformed(java.awt.event.ActionEvent evt) {
                refreshButtonActionPerformed(evt);
            }
        });

        javax.swing.GroupLayout layout = new javax.swing.GroupLayout(getContentPane());
        getContentPane().setLayout(layout);
        layout.setHorizontalGroup(
            layout.createParallelGroup(javax.swing.GroupLayout.Alignment.LEADING)
            .addGroup(layout.createSequentialGroup()
                .addContainerGap()
                .addGroup(layout.createParallelGroup(javax.swing.GroupLayout.Alignment.LEADING)
                    .addComponent(jScrollPane1, javax.swing.GroupLayout.DEFAULT_SIZE, 807, Short.MAX_VALUE)
                    .addGroup(javax.swing.GroupLayout.Alignment.TRAILING, layout.createSequentialGroup()
                        .addComponent(refreshButton)
                        .addPreferredGap(javax.swing.LayoutStyle.ComponentPlacement.RELATED, 588, Short.MAX_VALUE)
                        .addComponent(killButton)
                        .addGap(18, 18, 18)
                        .addComponent(migrateButton)))
                .addContainerGap())
        );
        layout.setVerticalGroup(
            layout.createParallelGroup(javax.swing.GroupLayout.Alignment.LEADING)
            .addGroup(javax.swing.GroupLayout.Alignment.TRAILING, layout.createSequentialGroup()
                .addContainerGap()
                .addComponent(jScrollPane1, javax.swing.GroupLayout.DEFAULT_SIZE, 602, Short.MAX_VALUE)
                .addPreferredGap(javax.swing.LayoutStyle.ComponentPlacement.RELATED)
                .addGroup(layout.createParallelGroup(javax.swing.GroupLayout.Alignment.BASELINE)
                    .addComponent(migrateButton)
                    .addComponent(killButton)
                    .addComponent(refreshButton))
                .addContainerGap())
        );

        pack();
    }// </editor-fold>//GEN-END:initComponents

	private void formWindowClosing(java.awt.event.WindowEvent evt) {//GEN-FIRST:event_formWindowClosing
		lock.unlock();
	}//GEN-LAST:event_formWindowClosing

	private void formWindowOpened(java.awt.event.WindowEvent evt) {//GEN-FIRST:event_formWindowOpened
		lock.lock();
		listProcs();
	}//GEN-LAST:event_formWindowOpened

	private void migrateButtonActionPerformed(java.awt.event.ActionEvent evt) {//GEN-FIRST:event_migrateButtonActionPerformed
		try {
			rpcConn.execute("session.meterpreter_write", new Object[]{session.get("id"),
					Base64.encode(("migrate "+processTable.getModel().getValueAt(processTable.getSelectedRow(),0) + "\n").getBytes())});
		} catch (Exception ex) {
			JOptionPane.showMessageDialog(null, ex);
		}
		listProcs();
	}//GEN-LAST:event_migrateButtonActionPerformed

	private void killButtonActionPerformed(java.awt.event.ActionEvent evt) {//GEN-FIRST:event_killButtonActionPerformed
		try {
			rpcConn.execute("session.meterpreter_write", new Object[]{session.get("id"),
					Base64.encode(("kill "+processTable.getModel().getValueAt(processTable.getSelectedRow(),0) + "\n").getBytes())});
		} catch (Exception ex) {
			JOptionPane.showMessageDialog(null, ex);
		}
		listProcs();
	}//GEN-LAST:event_killButtonActionPerformed

	private void refreshButtonActionPerformed(java.awt.event.ActionEvent evt) {//GEN-FIRST:event_refreshButtonActionPerformed
		listProcs();
	}//GEN-LAST:event_refreshButtonActionPerformed

    // Variables declaration - do not modify//GEN-BEGIN:variables
    private javax.swing.JScrollPane jScrollPane1;
    private javax.swing.JButton killButton;
    private javax.swing.JButton migrateButton;
    private javax.swing.JTable processTable;
    private javax.swing.JButton refreshButton;
    // End of variables declaration//GEN-END:variables
}
